#!env python
#-*- coding: utf-8 -*- 


def getTextFromSpeech(audio_src, model=None):
    """
    把输入的声音文件转为文本输出

    Example:
    ```python
    text = getTextFromSpeech(wav_file_path)
    ```
    """
    if model is None:
        from funasr import AutoModel
        model = AutoModel(
                model="paraformer-zh", model_revision="v2.0.4",
                vad_model="fsmn-vad", vad_model_revision="v2.0.4",
                punc_model="ct-punc-c", punc_model_revision="v2.0.4",
                disable_log=True,
                disable_pbar=False,
                use_timestamp=True
                )
    res = model.generate(input=audio_src)
    text = res[0]['text']
    return text, model


def generateTextFile(wav_file, text_file=None, model=None):
    """ 把 wav_file 的转换结果写入 text_file (默认值为把 wav_file 的扩展名换 成.txt )
    """
    from pathlib import Path
    assert(isinstance(wav_file, Path))
    if text_file is None:
        text_file = wav_file.parent / f'{wav_file.name[:-4]}.txt'
    print(f'convert {wav_file.name} to {text_file.name}')
    text, model = getTextFromSpeech(wav_file.as_posix(), model=model)
    with open(text_file, 'w') as file:
        print(text, file=file, flush=True)
    return text_file, model


def main(audio_src):
    from pathlib import Path
    audio_src_path = Path(audio_src)
    model = None
    if audio_src_path.exists():
        if audio_src_path.is_dir():
            for wav_file in audio_src_path.glob('*.wav'):
                text_file, model = generateTextFile(wav_file, model=model)
        else:
            assert(audio_src_path.is_file())
            wav_file = audio_src_path
            text_file, model = generateTextFile(wav_file, model=model)



if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("audio_src", help="path for input .wav file")
    args = parser.parse_args()
    main(args.audio_src)
